from colorama import Fore
from aalpy.base import SUL
from learning.CacheTree import CacheTree
from learning.DBhelper import *
from learning.Errors import NonDeterministicError
import utils
from logger import GLOBAL_LOGGER as LOG


###
# The code used in this file is copied from the SweynTooth project:
# https://github.com/DES-Lab/AALpy
#
# Following file/class has been copied:
# -- aalpy/base/SUL.py
#
# Adaptions to the existing code have been made:
# -- check for non-determinism
#
#
###
        
class FailSafeCacheSUL(SUL):
    """
    System under learning that keeps a multiset of all queries in memory.
    This multiset/cache is encoded as a tree.
    """
    def __init__(self, sul: SUL, database, must_query_result=None):
        super().__init__()
        self.sul = sul
        self.cache = CacheTree()
        self.dbhelper = DBhelper(database)
        self.must_query_result = must_query_result

    def clean_cache(self):
        self.cache = CacheTree()
        
    def query(self, word):
        """
        Performs a membership query on the SUL if and only if `word` is not a prefix of any trace in the cache.
        Before the query, pre() method is called and after the query post()
        method is called. Each letter in the word (input in the input sequence) is executed using the step method.

        Args:

            word: membership query (word consisting of letters/inputs)

        Returns:

            list of outputs, where the i-th output corresponds to the output of the system after the i-th input

        """
        LOG.logger.info('*'*100)
        LOG.logger.info('current query : ' + str(word))
        
        results_in_db_str = self.dbhelper.execute_query(str(word))
        if results_in_db_str:
            LOG.logger.info('found in sqlite : ' + results_in_db_str)
            results_in_db = str2list(results_in_db_str)
            self.cache.add_to_cache(word, results_in_db)
            return results_in_db
        
        cached_query = self.cache.in_cache(word)
        if cached_query:
            LOG.logger.info('found in cache tree : ' + str(cached_query))
            return cached_query

        repeated_times = 0
        while True:
            
            self.pre()
            try:
                out = [self.step(letter) for letter in word]
            except NonDeterministicError as e:
                LOG.logger.warning(Fore.RED + e.message)
                LOG.logger.warning(f"\033[31m{'#'*80}\033[0m")
                LOG.logger.warning('try to find most common respose')
                querys = e.querys
                expected_outputs = e.expected_outputs
                most_common_respose = get_most_respose(self.sul, querys)
                if expected_outputs != most_common_respose:
                    LOG.logger.warning('try to delete wrong respose in db')
                    self.dbhelper.clean_db(querys,most_common_respose)
                    LOG.logger.warning('try to clean cache tree')
                    self.clean_cache()
                self.post()
                continue
            
            # must_query_result是测试者认为必须满足的查询结果，可在配置文件中人为设置；
            # 如果目标服务器一直没有给出满足要求的结果，则认为目标服务器出现异常。
            pass_must_query = True
            if self.must_query_result != None:
                for expected_word, expected_result in self.must_query_result.items():
                    expected_word = utils.str2list(expected_word)
                    expected_result = utils.str2list(expected_result)
                    for prelen in range(1, len(expected_word)+1):
                        excepted_word2 = expected_word[:prelen]
                        excepted_result2 = expected_result[:prelen]
                        if list(word[:prelen]) == excepted_word2:
                            if out[:prelen] != excepted_result2:
                                repeated_times+=1
                                pass_must_query = False
                                if repeated_times >= 10:
                                    LOG.logger.error(Fore.RED + f"word: {word}")
                                    LOG.logger.error(Fore.RED + f"out: {out}")
                                    LOG.logger.error(Fore.RED + f"expected_word: {excepted_word2}")
                                    LOG.logger.error(Fore.RED + f"expected_result: {excepted_result2}")
                                    LOG.logger.error(Fore.RED + '----------------请检查设备活性----------------')
                                    self.post()
                                    time.sleep(300)
                        else:
                            break
                        
            self.post()
            if not pass_must_query:
                continue
            else:             
                break
        self.num_queries += 1
        self.num_steps += len(word)            
        
        self.cache.add_to_cache(word, out)
        self.dbhelper.insert_query(str(word), str(out))
        LOG.logger.info('insert to cachetree and database')
        return out


    def pre(self):
        """
        Reset the system under learning and current node in the cache tree.
        """
        self.cache.reset()
        self.sul.pre()

    def post(self):
        self.sul.post()

    def step(self, letter):
        """
        Executes an action on the system under learning, adds it to the cache and returns its result.

        Args:

           letter: Single input that is executed on the SUL.

        Returns:

           Output received after executing the input.

        """
        out = self.sul.step(letter)
        try:
            self.cache.step_in_cache(letter, out)
        except NonDeterministicError:
            raise
        
        return out


def get_most_respose(sul:SUL, word, max_test=7):
    resposes = []
    for _ in range(int(max_test)):
        resposes.append(str(sul.query(word)))
    return str2list(collections.Counter(resposes).most_common(1)[0][0])